from textual.widgets import DataTable
from textual.events import MouseEvent
from typing import TypeVar

from screens.dialog import DialogScreen

import time
import datetime
import psutil
import asyncio
import nest_asyncio

nest_asyncio.apply()
CellType = TypeVar("CellType")
interval = 2
cpus_usage = {}  # {pid: cpu_percent}
# 网络下载、上传数据量
now_in = psutil.net_io_counters().bytes_recv
now_out = psutil.net_io_counters().bytes_sent


class MyDataTable(DataTable):
    
    def __init__(self, id):
        super().__init__(id=id)
        self.interval = interval
        self.p_cpu_per = 0.0

    def kill_process(self, pid):
        try:
            parent_proc = psutil.Process(pid)
            for child_proc in parent_proc.children(recursive=True):
                child_proc.kill()
            parent_proc.kill()
        except:
            pass

    def update_pids_label(self):
        global now_in, now_out
        delta_in = psutil.net_io_counters().bytes_recv - now_in
        delta_out = psutil.net_io_counters().bytes_sent - now_out

        def get_delta_unit(delta):
            a = delta/1024
            b = a/1024
            c = b/1024
            delta_unit = 'B/s'
            if round(c) != 0:
                delta = round(c/interval, 2)
                delta_unit = 'GB/s'
            elif round(b) != 0:
                delta = round(b/interval, 2)
                delta_unit = 'MB/s'
            elif round(a) != 0:
                delta = round(a/interval, 2)
                delta_unit = 'KB/s'
            return delta, delta_unit

        delta_in, delta_in_unit = get_delta_unit(delta_in)
        delta_out, delta_out_unit = get_delta_unit(delta_out)
        
        # 上传下载速率
        now_in = psutil.net_io_counters().bytes_recv
        now_out = psutil.net_io_counters().bytes_sent

        # 运行中进程数
        running = 0
        for i in psutil.process_iter():
            if i.status()=="running":
                running += 1

        self.border_title = "进程共" + str(len(psutil.pids()))+"个;运行中"+str(running)+";网络上传:"+str(delta_out)+delta_out_unit+";下载:"+str(delta_in)+delta_in_unit

    def key_space(self):
        if self.has_focus:
            row = self.get_row_at(self.cursor_row)
            pid = row[0]
            row_key = str(pid)
            pid_create_time = datetime.datetime.fromtimestamp(psutil.Process(pid).create_time()).strftime("%Y-%m-%d %H:%M:%S")
            self.app.install_screen(
                DialogScreen(
                    text=f"\n提示: 是否终止进程{row[0]}及其所有子进程?\n\n名称: {row[4]}\n状态: {psutil.Process(pid).status()}\nCPU(%): {row[2]}\n内存(%): {row[3]}\n创建时间: {pid_create_time}\n所有者: {row[5]}\n使用终端: {psutil.Process(pid).terminal()}\n可执行文件路径: {psutil.Process(pid).exe()}\n进程工作目录: {psutil.Process(pid).cwd()}"), 
                    name="dialog")
            def check_result(result: bool):
                if result:
                    self.kill_process(pid)
                    self.remove_row(str(row_key))
                self.app.uninstall_screen("dialog")
            self.app.push_screen("dialog", check_result)

    def on_click(self, event:MouseEvent):
        # 1: left   2:middle   3:right
        if event.button==3 and self.has_focus:
            row = self.get_row_at(self.cursor_row)
            pid = row[0]
            row_key = str(pid)
            pid_create_time = datetime.datetime.fromtimestamp(psutil.Process(pid).create_time()).strftime("%Y-%m-%d %H:%M:%S")
            self.app.install_screen(
                DialogScreen(
                    text=f"\n提示: 是否终止进程{row[0]}及其所有子进程?\n\n名称: {row[4]}\n状态: {psutil.Process(pid).status()}\nCPU(%): {row[2]}\n内存(%): {row[3]}\n创建时间: {pid_create_time}\n所有者: {row[5]}\n使用终端: {psutil.Process(pid).terminal()}\n可执行文件路径: {psutil.Process(pid).exe()}\n进程工作目录: {psutil.Process(pid).cwd()}"), 
                    name="dialog")
            def check_result(result: bool):
                if result:
                    self.kill_process(pid)
                    self.remove_row(str(row_key))
                self.app.uninstall_screen("dialog")
            self.app.push_screen("dialog", check_result)
        if event.button==1 and self.has_focus:
            self.app.query_one("#tree_box").styles.border = ("solid", "white")
            self.app.query_one("#user_table").styles.border = ("solid", "white")
            self.styles.border = ("solid", "orangered")
            self.app.query_one("#user_table").cursor_type = 'none'
            self.cursor_type = 'row'

    def key_tab(self):
        self.app.query_one("#tree_box").styles.border = ("solid", "orangered")
        self.styles.border = ("solid", "white")
        self.cursor_type = 'none'
    
    def get_pid_status(self, p):
        # status: 1 -> running; 0 -> sleeping; -1 ->idle.
        s = 0
        if p.status()=="running":
            s = 1
        elif p.status()=="sleeping":
            s = 0
        elif p.status()=="disk_sleep":
            s = -1
        elif p.status()=="stopped":
            s = -2
        elif p.status()=="tracing_stop":
            s = -3
        elif p.status()=="zombie":
            s = -4
        elif p.status()=="dead":
            s = -5
        elif p.status()=="wake_kill":
            s = -6
        elif p.status()=="waking":
            s = -7
        elif p.status()=="parked":
            s = -8
        elif p.status()=="idle":
            s = -3
        elif p.status()=="locked":
            s = -9
        elif p.status()=="waiting":
            s = -10
        elif p.status()=="suspended":
            s = -11
        return s

    def on_mount(self):
        self.header = ("pid", "状态", "CPU(%)", "内存(%)", "名称", "所有者", "父进程pid")
        self.cursor_type = "none"
        self.zebra_stripes = True
        self.pids_nums = len(psutil.pids())
        # self.table.add_columns(*self.header)
        for j in self.header:
            self.add_column(j, key=j)
        self.pids=psutil.pids()
        for i in range(len(self.pids)):
            p = psutil.Process(self.pids[i])
            s = self.get_pid_status(p)
            self.add_row(p.pid, 
                         s,
                         0.00,  # 初始0.00
                         round(p.memory_percent(),2),
                         p.name(),
                         p.username(), 
                         str([i.pid for i in p.parents()]).strip("[]"), 
                         key=str(self.pids[i]))
        self.sort("状态","CPU(%)","内存(%)", reverse=True)
        self.set_interval(self.interval, self.update_datas, pause=False)

    async def get_cpu_percent(self, pids, interval_1):
        global cpus_usage
        # 自定义获取cpu利用率
        num_cpus = psutil.cpu_count()
        _timer = getattr(time, 'monotonic', time.time)
        st1_list = []
        pt1_list = []
        st2_list = []
        pt2_list = []
        def timer():
            return _timer() * num_cpus
 
        # 记录第一次时间
        for i in pids:
            try:
                pi = psutil.Process(i)
            except:
                cpus_usage[i] = 0.00
                st1_list.append(None)
                pt1_list.append(None)
                continue
            with pi.oneshot():
                st1 = timer()
                pt1 = pi.cpu_times()
                st1_list.append(st1)
                pt1_list.append(pt1)
        # 等待
        await asyncio.sleep(interval_1)
        # 记录第二次时间
        for j in pids:
            try:
                pj = psutil.Process(j)
            except:
                cpus_usage[j] = 0.00
                st2_list.append(None)
                pt2_list.append(None)
                continue
            with pj.oneshot():
                st2 = timer()
                pt2 = pj.cpu_times()
                st2_list.append(st2)
                pt2_list.append(pt2)
        # 计算占用率
        for m in range(len(pids)):
            st1_1 = st1_list[m]
            pt1_1 = pt1_list[m]
            st2_2 = st2_list[m]
            pt2_2 = pt2_list[m]
            if st1_1 != None and pt1_1 != None and st2_2 != None and pt2_2 != None:
                delta_proc = (pt2_2.user - pt1_1.user) + (pt2_2.system - pt1_1.system)
                delta_time = st2_2 - st1_1
            else:
                delta_proc = 0.00
                delta_time = 0.00
            try:
                cpu_usage = ((delta_proc / delta_time) * 100) / num_cpus
                cpus_usage[pids[m]] = cpu_usage
            except:
                cpus_usage[pids[m]] = 0.00
  
    async def update_datas(self):
        self.update_pids_label()
        pids = psutil.pids()
        new_pids = list(set(pids) - set(self.pids))
        old_pids = []
        f = self.get_cpu_percent(pids, 1.5)
        asyncio.run(f)
        for i in range(len(self.rows)):
            # 获取行信息
            row = self.get_row_at(i)
            pid = row[0]

            # 储存已经消失的进程，最后一并删除
            if pid not in pids:
                old_pids.append(str(pid))
            # 更新其他进程状态
            try:
                p = psutil.Process(pid)
                s = self.get_pid_status(p)
                self.update_cell(str(pid),'状态',s)
                self.update_cell(str(pid),'名称',p.name())
                self.update_cell(str(pid),'所有者',p.username())
                self.update_cell(str(pid),'CPU(%)',round(cpus_usage[pid] if cpus_usage[pid]!=None else 0.00,2))
                self.update_cell(str(pid),'内存(%)',round(p.memory_percent(),2))
                self.update_cell(str(pid),'父进程pid',str([i.pid for i in p.parents()]).strip("[]"))
            except:
                pass

        for n in old_pids:
            self.remove_row(n)

        for j in new_pids:
            # 添加新的进程
            try:
                p1 = psutil.Process(j)
                s=self.get_pid_status(p1)
                self.add_row(p1.pid, 
                    s, 
                    0.00,  # 初始0.00
                    round(p1.memory_percent(),2), 
                    # p1.name().split('/')[0],
                    p1.name(),
                    p1.username(), 
                    str([i.pid for i in p1.parents()]).strip("[]"), 
                    key=str(p1.pid))
            except:
                pids.remove(j)
        self.pids = pids
        self.sort("状态","CPU(%)","内存(%)", reverse=True)


        






